package org.hong.monkey.rpc.netty

import java.io.File
import java.util.concurrent.ConcurrentHashMap

import org.hong.monkey.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
import org.hong.monkey.network.server.StreamManager
import org.hong.monkey.rpc.RpcEnvFileServer
import org.hong.monkey.util.Utils

/**
 * StreamManager implementation for serving files from a NettyRpcEnv.
 */
private[netty] class NettyStreamManager(rpcEnv: NettyRpcEnv)
  extends StreamManager with RpcEnvFileServer {

  private val files = new ConcurrentHashMap[String, File]()
  private val jars = new ConcurrentHashMap[String, File]()

  override def getChunk(streamId: Long, chunkIndex: Int): ManagedBuffer = {
    throw new UnsupportedOperationException()
  }

  override def openStream(streamId: String): ManagedBuffer = {
    val Array(ftype, fname) = streamId.stripPrefix("/").split("/", 2)
    val file = ftype match {
      case "files" => files.get(fname)
      case "jars" => jars.get(fname)
      case _ => throw new IllegalArgumentException(s"Invalid file type: $ftype")
    }

    require(file != null && file.isFile(), s"File not found: $streamId")
    new FileSegmentManagedBuffer(rpcEnv.transportConf, file, 0, file.length())
  }

  override def addFile(file: File): String = {
    require(files.putIfAbsent(file.getName(), file) == null,
      s"File ${file.getName()} already registered.")
    s"${rpcEnv.address.toMonkeyURL}/files/${Utils.encodeFileNameToURIRawPath(file.getName())}"
  }

  override def addJar(file: File): String = {
    require(jars.putIfAbsent(file.getName(), file) == null,
      s"JAR ${file.getName()} already registered.")
    s"${rpcEnv.address.toMonkeyURL}/jars/${Utils.encodeFileNameToURIRawPath(file.getName())}"
  }

}
